from torch.utils import data
import torch
from torchvision import transforms
import os
import numpy as np

from osgeo import gdal

def _pre(img, min, max):
    img = ((img - min) / (max - min)).astype(np.float32)
    if img.ndim == 2:
        img = np.stack([img,]*3, axis=-1)
    return img


class gdal_dataset(data.Dataset):
    def __init__(self, img, patch_size, oversize, sf):

        self.oversize = oversize  # overlap of two image
        self.GPU_patch_size = patch_size  # how large the GPU process
        self.patch_size = self.GPU_patch_size - self.oversize  # reading interval
        self.sf = sf

        self.img_L = gdal.Open(img)
        self.img_width = self.img_L.RasterXSize
        self.img_height = self.img_L.RasterYSize
        img_type = self.img_L.ReadAsArray(0, 0, 1, 1).dtype.name
        self.bandsnum = self.img_L.RasterCount
        print("bandsnum: ", self.bandsnum)
        print("img_type:", img_type)
        self.maxvalue = 255
        self.minvalue = 0
        if img_type != 'uint8':
            if self.bandsnum == 3:
                maxmin = np.zeros((3, 2))
                maxmin[0] = self.img_L.GetRasterBand(1).ComputeRasterMinMax()
                maxmin[1] = self.img_L.GetRasterBand(2).ComputeRasterMinMax()
                maxmin[2] = self.img_L.GetRasterBand(3).ComputeRasterMinMax()
                self.maxvalue = max(maxmin[:, 1])
                self.minvalue = min(maxmin[:, 0])
            else:
                maxmin = self.img_L.GetRasterBand(1).ComputeRasterMinMax()
                self.maxvalue = maxmin[1]
                self.minvalue = maxmin[0]


        self.range_y = self.img_height // self.patch_size + 1
        self.range_x = self.img_width // self.patch_size + 1

        self.waiting_list = []
        for i in range(self.range_y):
            for j in range(self.range_x):
                self.waiting_list.append([i, j])
        print("range_y", self.range_y)
        print("range_x", self.range_x)
        print("height, width", self.img_height, self.img_width)
        self.data_transformer = transforms.Compose([
            transforms.Lambda(lambda img : _pre(img, self.minvalue, self.maxvalue)),
            transforms.ToTensor()
        ])

    def __getitem__(self, item):
        i, j = self.waiting_list[item]

        readj = j * self.patch_size
        readi = i * self.patch_size
        if self.GPU_patch_size + readi > self.img_height:
            readi = self.img_height - self.GPU_patch_size
        if self.GPU_patch_size + readj > self.img_width:
            readj = self.img_width - self.GPU_patch_size

        img_data = self.img_L.ReadAsArray(int(readj), int(readi), self.GPU_patch_size, self.GPU_patch_size)
        if self.bandsnum == 3:
            return torch.from_numpy(np.ascontiguousarray(img_data)).float().div(255.)
        else:
            return self.data_transformer(img_data)

    def __len__(self):
        return self.range_x * self.range_y




if __name__ == '__main__':
    import matplotlib.pyplot as plt
    img_path = "testsets/2/GF02_MS1_038712_20211017_KS450_01_003_L1A_011.tif"
    patch_size = 1000
    oversize = 100
    dataset = gdal_dataset(img_path, patch_size, oversize, 2)
    dataloader = data.DataLoader(dataset, batch_size=2, shuffle=False, num_workers=2)
    for item, sample in enumerate(dataloader):
        img_batch = sample.numpy()
        print(img_batch.shape)
        for jj in range(sample.size()[0]):
            img_tmp = np.transpose(img_batch[jj], axes=[1, 2, 0])
            print(img_tmp.shape)
            img_tmp = img_tmp * 255.0
            plt.figure()
            plt.title('display')
            plt.imshow(img_tmp.astype(np.uint8))
            plt.show(block=True)



